import tkinter as tk
from tkinter import filedialog, font
from PIL import ImageGrab, ImageTk
from datetime import datetime
import win32gui
import os


class ScreenshotTool:
    def __init__(self, root):
        self.root = root
        self.root.title("截图工具")
        self.root.geometry("80x40")
        self.root.resizable(False, False)
        self.root.attributes('-toolwindow', True, '-topmost', True)
        self._create_screenshot_button()
        self.callback = None

    def set_callback(self, callback):
        self.callback = callback

    def _create_screenshot_button(self):
        button_font = font.Font(family="微软雅黑", size=10)
        self.screenshot_button = tk.Button(self.root, text="截图", command=self.start_screenshot, font=button_font)
        self.screenshot_button.pack(pady=0, expand=True, fill=tk.BOTH)

    def start_screenshot(self):
        self.root.withdraw()
        screen_width = self.root.winfo_screenwidth()
        screen_height = self.root.winfo_screenheight()
        self._create_screenshot_window()
        self.start_x = None
        self.start_y = None
        self.rect = None

    def _create_screenshot_window(self):
        self.screenshot_window = tk.Toplevel()
        self.screenshot_window.attributes('-fullscreen', True)
        self.screenshot_window.attributes('-alpha', 0.3)
        self.screenshot_window.configure(bg='black')
        self.screenshot_window.resizable(False, False)
        self.screenshot_window.attributes('-toolwindow', True)
        self.canvas = tk.Canvas(self.screenshot_window, bg='black', highlightthickness=0)
        self.canvas.pack(fill=tk.BOTH, expand=True)
        self.canvas.bind("<ButtonPress-1>", self.on_press)
        self.canvas.bind("<B1-Motion>", self.on_drag)
        self.canvas.bind("<ButtonRelease-1>", self.on_release)

    def on_press(self, event):
        self.start_x = event.x
        self.start_y = event.y
        if self.rect:
            self.canvas.delete(self.rect)

    def on_drag(self, event):
        if self.start_x and self.start_y:
            if self.rect:
                self.canvas.delete(self.rect)
            self.rect = self.canvas.create_rectangle(self.start_x, self.start_y, event.x, event.y, outline='red', width=2)

    def on_release(self, event):
        end_x = event.x
        end_y = event.y
        left = min(self.start_x, end_x)
        top = min(self.start_y, end_y)
        right = max(self.start_x, end_x)
        bottom = max(self.start_y, end_y)
        self.screenshot_window.destroy()
        self.root.deiconify()
        screenshot = ImageGrab.grab(bbox=(left, top, right, bottom))
        self._create_confirm_window(screenshot)

    def _create_confirm_window(self, screenshot):
        self.confirm_window = tk.Toplevel(self.root)
        self.confirm_window.title("确认保存")
        self.confirm_window.resizable(False, False)
        self.confirm_window.attributes('-toolwindow', True)
        self.confirm_window.attributes('-topmost', True)

        photo = ImageTk.PhotoImage(screenshot)
        label = tk.Label(self.confirm_window, image=photo)
        label.image = photo
        label.pack(pady=0)
        confirm_button = tk.Button(self.confirm_window, text="确认保存", command=lambda: self.save_screenshot(screenshot))
        confirm_button.pack(side=tk.LEFT, padx=2, pady=2)
        cancel_button = tk.Button(self.confirm_window, text="取消", command=self.confirm_window.destroy)
        cancel_button.pack(side=tk.RIGHT, padx=2, pady=2)
        self.confirm_window.update_idletasks()
        screen_width = self.root.winfo_screenwidth()
        screen_height = self.root.winfo_screenheight()
        window_width = self.confirm_window.winfo_width()
        window_height = self.confirm_window.winfo_height()
        x = (screen_width - window_width) // 2
        y = (screen_height - window_height) // 2
        self.confirm_window.geometry(f"+{x}+{y}")

    def save_screenshot(self, screenshot):
        try:
            save_dir = filedialog.askdirectory(title="选择保存目录")
            if not save_dir:
                return
            current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"{save_dir}/screenshot_{current_time}.png"
            screenshot.save(filename)
            self.show_save_message(f"截图已保存为 {filename}")
            if self.callback:
                self.callback(filename)
        except Exception as e:
            self.show_error_message(f"保存截图时出错: {e}")
        finally:
            self.confirm_window.destroy()

    def show_save_message(self, message):
        message_window = tk.Toplevel(self.root)
        message_window.title("提示")
        message_window.attributes('-toolwindow', True)
        message_window.resizable(False, False)
        label = tk.Label(message_window, text=message)
        label.pack(pady=10)
        message_window.update_idletasks()
        screen_width = self.root.winfo_screenwidth()
        window_width = message_window.winfo_width()
        x = (screen_width - window_width) // 2
        y = 20
        message_window.geometry(f"+{x}+{y}")
        self.root.after(2000, message_window.destroy)

    def show_error_message(self, message):
        message_window = tk.Toplevel(self.root)
        message_window.title("错误")
        message_window.attributes('-toolwindow', True)
        message_window.resizable(False, False)
        label = tk.Label(message_window, text=message, fg='red')
        label.pack(pady=10)
        message_window.update_idletasks()
        screen_width = self.root.winfo_screenwidth()
        window_width = message_window.winfo_width()
        x = (screen_width - window_width) // 2
        y = 20
        message_window.geometry(f"+{x}+{y}")
        self.root.after(3000, message_window.destroy)

def check_main_program():
    main_program_running = False
    def enum_windows_callback(hwnd, _):
        nonlocal main_program_running
        window_title = win32gui.GetWindowText(hwnd)
        if window_title.startswith("识图助手V"):
            main_program_running = True

    win32gui.EnumWindows(enum_windows_callback, None)
    if not main_program_running:
        error_window = tk.Toplevel()
        error_window.title("错误")
        error_window.attributes('-toolwindow', True)
        error_window.resizable(False, False)
        error_font = font.Font(family="微软雅黑", size=24, weight="bold")
        label = tk.Label(error_window, text="未启动主程序", fg='red', font=error_font)
        label.pack(pady=10)
        error_window.update_idletasks()
        screen_width = root.winfo_screenwidth()
        screen_height = root.winfo_screenheight()
        window_width = error_window.winfo_width()
        window_height = error_window.winfo_height()
        x = (screen_width - window_width) // 2
        y = (screen_height - window_height) // 2
        error_window.geometry(f"+{x}+{y}")
        root.after(2000, lambda: force_exit())
        return False
    return True

def force_exit():
    os._exit(1)

if __name__ == "__main__":
    root = tk.Tk()
    if check_main_program():
        app = ScreenshotTool(root)
    root.mainloop()